{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Apply OpenVINO Acceleration on Inference Pipeline"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Environment Preparation\n",
    "Before you start with APIs delivered by bigdl-nano, you have to make sure BigDL-Nano is correctly installed for PyTorch. If not, please follow [this](../../../../../docs/readthedocs/source/doc/Nano/Overview/nano.md) to set up your environment.<br><br>\n",
    "To use OpenVINO acceleration, you have to install the OpenVINO toolkit:\n",
    "```bash\n",
    "pip install openvino-dev\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 0: Load Data\n",
    "We used two images from [Dogs vs. Cats datasets](https://www.kaggle.com/c/dogs-vs-cats) for demo."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import torch\n",
    "from torchvision.io import read_image\n",
    "from torchvision import transforms\n",
    "from torchvision.datasets import OxfordIIITPet\n",
    "from torch.utils.data.dataloader import DataLoader\n",
    "\n",
    "train_transform = transforms.Compose([transforms.Resize(256),\n",
    "                                      transforms.RandomCrop(224),\n",
    "                                      transforms.RandomHorizontalFlip(),\n",
    "                                      transforms.ColorJitter(brightness=.5, hue=.3),\n",
    "                                      transforms.ToTensor(),\n",
    "                                      transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])])\n",
    "val_transform = transforms.Compose([transforms.Resize([224, 224]), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])])\n",
    "# Apply data augmentation to the tarin_dataset\n",
    "train_dataset = OxfordIIITPet(root = \"/tmp/data\", transform=train_transform, download=True)\n",
    "val_dataset = OxfordIIITPet(root=\"/tmp/data\", transform=val_transform)\n",
    "\n",
    "# obtain training indices that will be used for validation\n",
    "indices = torch.randperm(len(train_dataset))\n",
    "val_size = len(train_dataset) // 4\n",
    "train_dataset = torch.utils.data.Subset(train_dataset, indices[:-val_size])\n",
    "val_dataset = torch.utils.data.Subset(val_dataset, indices[-val_size:])\n",
    "\n",
    "# prepare data loaders\n",
    "train_dataloader = DataLoader(train_dataset, batch_size=32)\n",
    "val_dataloader = DataLoader(val_dataset, batch_size=32)\n",
    "\n",
    "DEV_RUN = bool(os.environ.get('DEV_RUN', False))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 2: Custom Model\n",
    "Regarding the model, we used pretrained torchvision.models.resnet18. More details, please refer to [here](https://pytorch.org/vision/0.12/generated/torchvision.models.resnet18.html?highlight=resnet18)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "GPU available: False, used: False\n",
      "TPU available: False, using: 0 TPU cores\n",
      "IPU available: False, using: 0 IPUs\n",
      "/opt/conda/envs/testOpenvino/lib/python3.7/site-packages/pytorch_lightning/trainer/trainer.py:532: LightningDeprecationWarning: `trainer.fit(train_dataloader)` is deprecated in v1.4 and will be removed in v1.6. Use `trainer.fit(train_dataloaders)` instead. HINT: added 's'\n",
      "  \"`trainer.fit(train_dataloader)` is deprecated in v1.4 and will be removed in v1.6.\"\n",
      "\n",
      "  | Name  | Type             | Params\n",
      "-------------------------------------------\n",
      "0 | model | ResNet           | 11.2 M\n",
      "1 | loss  | CrossEntropyLoss | 0     \n",
      "-------------------------------------------\n",
      "11.2 M    Trainable params\n",
      "0         Non-trainable params\n",
      "11.2 M    Total params\n",
      "44.782    Total estimated model params size (MB)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Validation sanity check:   0%|          | 0/2 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/opt/conda/envs/testOpenvino/lib/python3.7/site-packages/pytorch_lightning/trainer/data_loading.py:106: UserWarning: The dataloader, val dataloader 0, does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` (try 96 which is the number of cpus on this machine) in the `DataLoader` init to improve performance.\n",
      "  f\"The dataloader, {name}, does not have many workers which may be a bottleneck.\"\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "                                                                      "
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/opt/conda/envs/testOpenvino/lib/python3.7/site-packages/pytorch_lightning/trainer/data_loading.py:106: UserWarning: The dataloader, train dataloader, does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` (try 96 which is the number of cpus on this machine) in the `DataLoader` init to improve performance.\n",
      "  f\"The dataloader, {name}, does not have many workers which may be a bottleneck.\"\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 4: 100%|██████████| 116/116 [00:49<00:00,  2.37it/s, loss=0.286, v_num=14, val/loss=1.030] \n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "tensor([11,  4])"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import torch\n",
    "from torchvision.models import resnet18\n",
    "from bigdl.nano.pytorch import Trainer\n",
    "\n",
    "model_ft = resnet18(pretrained=True)\n",
    "num_ftrs = model_ft.fc.in_features\n",
    "\n",
    "# Here the size of each output sample is set to 37.\n",
    "model_ft.fc = torch.nn.Linear(num_ftrs, 37)\n",
    "loss_ft = torch.nn.CrossEntropyLoss()\n",
    "optimizer_ft = torch.optim.SGD(model_ft.parameters(), lr=0.01, momentum=0.9, weight_decay=5e-4)\n",
    "\n",
    "# Compile our model with loss function, optimizer.\n",
    "model = Trainer.compile(model_ft, loss_ft, optimizer_ft)\n",
    "trainer = Trainer(max_epochs=5,\n",
    "                  fast_dev_run=DEV_RUN) # Run model quickly in test\n",
    "trainer.fit(model, train_dataloaders=train_dataloader)\n",
    "\n",
    "# Inference/Prediction\n",
    "x = torch.stack([val_dataset[0][0], val_dataset[1][0]])\n",
    "model_ft.eval()\n",
    "y_hat = model_ft(x)\n",
    "y_hat.argmax(dim=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### OpenVINO Acceleration\n",
    "trace your model as a openvino model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model Optimizer arguments:\n",
      "Common parameters:\n",
      "\t- Path to the Input Model: \t/tmp/tmpddetxksv/tmp.onnx\n",
      "\t- Path for generated IR: \t/tmp/tmpiafile6h\n",
      "\t- IR output name: \ttmp\n",
      "\t- Log level: \tERROR\n",
      "\t- Batch: \tNot specified, inherited from the model\n",
      "\t- Input layers: \tNot specified, inherited from the model\n",
      "\t- Output layers: \tNot specified, inherited from the model\n",
      "\t- Input shapes: \tNot specified, inherited from the model\n",
      "\t- Source layout: \tNot specified\n",
      "\t- Target layout: \tNot specified\n",
      "\t- Layout: \tNot specified\n",
      "\t- Mean values: \tNot specified\n",
      "\t- Scale values: \tNot specified\n",
      "\t- Scale factor: \tNot specified\n",
      "\t- Precision of IR: \tFP32\n",
      "\t- Enable fusing: \tTrue\n",
      "\t- User transformations: \tNot specified\n",
      "\t- Reverse input channels: \tFalse\n",
      "\t- Enable IR generation for fixed input shape: \tFalse\n",
      "\t- Use the transformations config file: \tNone\n",
      "Advanced parameters:\n",
      "\t- Force the usage of legacy Frontend of Model Optimizer for model conversion into IR: \tFalse\n",
      "\t- Force the usage of new Frontend of Model Optimizer for model conversion into IR: \tFalse\n",
      "OpenVINO runtime found in: \t/opt/conda/envs/testOpenvino/lib/python3.7/site-packages/openvino\n",
      "OpenVINO runtime version: \t2022.1.0-7019-cdb9bec7210-releases/2022/1\n",
      "Model Optimizer version: \t2022.1.0-7019-cdb9bec7210-releases/2022/1\n",
      "[ SUCCESS ] Generated IR version 11 model.\n",
      "[ SUCCESS ] XML file: /tmp/tmpiafile6h/tmp.xml\n",
      "[ SUCCESS ] BIN file: /tmp/tmpiafile6h/tmp.bin\n",
      "[ SUCCESS ] Total execution time: 0.50 seconds. \n",
      "[ SUCCESS ] Memory consumed: 183 MB. \n",
      "It's been a while, check for a new version of Intel(R) Distribution of OpenVINO(TM) toolkit here https://software.intel.com/content/www/us/en/develop/tools/openvino-toolkit/download.html?cid=other&source=prod&campid=ww_2022_bu_IOTG_OpenVINO-2022-1&content=upg_all&medium=organic or on the GitHub*\n",
      "[ INFO ] The model was converted to IR v11, the latest model format that corresponds to the source DL framework input/output format. While IR v11 is backwards compatible with OpenVINO Inference Engine API v1.0, please use API v2.0 (as of 2022.1) to take advantage of the latest improvements in IR v11.\n",
      "Find more information about API v2.0 and IR v11 at https://docs.openvino.ai\n"
     ]
    }
   ],
   "source": [
    "from bigdl.nano.pytorch import InferenceOptimizer\n",
    "ov_model = InferenceOptimizer.trace(model_ft, accelerator=\"openvino\", input_sample=torch.rand(1, 3, 224, 224))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The usage is almost the same with any PyTorch module"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([11,  4])"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "y_hat = ov_model(x)\n",
    "y_hat.argmax(dim=1)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.7.10 ('testOpenvino')",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.10"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "edf698c9356c4f0740ac364d5e27207f5ab66fab651ea6afd333d7f1671f1a40"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
